'''
Created on Dec 9, 2009

@author: mkiyer
'''
import gzip
import sys
import numpy as np
from veggie.genome.chrom import get_genome

class WiggleFileWriter(object):

    WRITING = 1
    SEARCHING = 0
            
    def __init__(self, filename=None, compress=True, span=-1):
        if filename is None:
            self.fhd = sys.stdout
            self.needs_close = False
        elif compress is True:
            self.fhd = gzip.open(filename, 'wb', compresslevel=5)
            self.needs_close = True
        else:
            self.fhd = open(filename, 'w')
            self.needs_close = True
        self.chrom = None
        self.pos = None
        self.state = self.SEARCHING
        # specific to write-span
        self.span = span
        if span > 0:
            self.span = span
            self.buffer = None
            self.mod_end = None

    def close(self):
        if self.span > 0:
            self.flush_span()
        if self.needs_close:
            self.fhd.close()
    
    def write_span(self, chrom, start, end, arr):
        # if start is not a multiple of span then try to retrieve
        # previous data from buffer
        mod_start = start - (start % self.span)
        if mod_start != start:
            if ((self.buffer is not None) and (len(self.buffer) == (start - mod_start))):
                mod_arr = np.append(self.buffer, arr)
            else:
                # cannot retrieve previous data so add zeros
                mod_arr = np.append(np.zeros(start - mod_start), arr)
        else:
            mod_arr = arr
        # if end not a multiple of span then save the extra data
        mod_end = end - (end % self.span)
        self.mod_end = mod_end
        if mod_end != end:
            self.buffer = arr[mod_end - start:end - start]
        # make sure we are writing to the correct chromosome
        if self.chrom != chrom:            
            self.fhd.write("variableStep chrom=%s span=%d\n" % (chrom, self.span))
            self.chrom = chrom
        # write the values
        for pos in xrange(mod_start, mod_end, self.span):
            segment = np.array(mod_arr[pos - mod_start:pos - mod_start + self.span])
            if segment.any():
                self.fhd.write("%d\t%.2f\n" % (pos + 1, np.average(segment)))

    def flush_span(self):
        if ((self.buffer is not None) and (len(self.buffer) > 0)):            
            segment = np.append(self.buffer, np.zeros(self.span - len(self.buffer)))
            if segment.any():
                self.fhd.write("%d\t%.2f\n" % (self.mod_end + 1, np.average(segment)))            
    
    def write_variable_step(self, chrom, start, end, arr, span=1):
        if self.chrom != chrom:            
            self.fhd.write("variableStep chrom=%s span=%d\n" % (chrom, span))
            self.chrom = chrom          
        indexes = np.flatnonzero(arr)
        lines = [("%d\t%d\n" % (start + pos + 1, cov)) for pos, cov in zip(indexes, arr[indexes])]
        self.fhd.writelines(lines)

    def write_fixed_step(self, chrom, start, end, arr, span=-1):
        # currently not writing at a specific location, don't need to do anything
        if self.state == self.SEARCHING:
            if self.chrom is None:
                self.chrom = chrom
            if self.pos is None:
                self.pos = start
        elif self.state == self.WRITING:
            # if previously writing to file, either need to flush
            # or continue writing across boundary
            # check position of this write against current position
            if not ((self.chrom == chrom) and (self.pos == start - 1)):
                # restart write at new position
                self.chrom = chrom
                self.pos = start
                self.state = self.SEARCHING
        
        # buffer and do block writes
        for i in xrange(len(arr)):
            self.pos = start + i
            cov = arr[i]            
            if cov != 0:
                if self.state == self.SEARCHING:
                    # write leading zero so that y axis scale includes zero
#                    if self.pos > 0:
#                        self.fhd.write("fixedStep chrom=%s start=%d step=1\n" % (self.chrom, self.pos))
#                        self.fhd.write("0\n")
#                    else:
                    self.fhd.write("fixedStep chrom=%s start=%d step=1\n" % (self.chrom, self.pos+1)) 
                    self.state = self.WRITING            
                self.fhd.write("%d\n" % cov)
            else:
                self.state = self.SEARCHING
                # write trailing zero so that y axis scale includes zero
    #            if state == WRITING:
    #                if pos < end:
    #                    wigglefile.write("0\n")
        return
    
class WiggleFileReader(object):
    '''
    Reader for wiggle (WIG) files
    
    https://cgwb.nci.nih.gov/goldenPath/help/wiggle.html
     
    track type=wiggle_0 name=track_label description=center_label \
        visibility=display_mode color=r,g,b altColor=r,g,b \
        priority=priority autoScale=on|off \
        gridDefault=on|off maxHeightPixels=max:default:min \
        graphType=bar|points viewLimits=lower:upper \
        yLineMark=real-value yLineOnOff=on|off \
        windowingFunction=maximum|mean|minimum smoothingWindow=off|2-16
    '''
    # parser modes
    TRACK = 0
    VARIABLE_STEP = 1
    FIXED_STEP = 2
    BED = 3

    def __init__(self, line_iterator, dtype):
        self.line_iter = line_iterator
        self.state = None
        self.chrom = None
        self.start = None
        self.pos = None
        self.step = 1
        self.span = 1
        self.dtype = dtype
            
    def _parse_fixed_step_header(self, line):
        # fixedStep  chrom=chrN  start=position  step=stepInterval  [span=windowSize]
        chrom, start, step, span = None, -1, 1, 1        
        fields = line.split()[1:]
        for field in fields:
            k,v = field.split('=', 1)
            if k == 'chrom':
                chrom = v
            elif k == 'start':
                start = int(v)
            elif k == 'step':
                step = int(v)
            elif k == 'span':
                span = int(v)
        return chrom, start, step, span 
    
    def _parse_variable_step_header(self, line):
        # variableStep  chrom=chrN  [span=windowSize]
        chrom, span = None, 1
        fields = line.split()[1:]
        for field in fields:
            k,v = field.split('=', 1)
            if k == 'chrom':
                chrom = v
            elif k == 'span':
                span = int(v)
        return chrom, span
                            
    def iterchunks(self, chunk_size=(2**20)):
        import numpy as np        

        arr = np.zeros(chunk_size, dtype=self.dtype)
        self.start = None
        self.pos = None
        
        for line in self.line_iter:
            line = line.strip()
            if line.startswith('track'):
                pass
            elif line.startswith('fixedStep'):
                self.state = self.FIXED_STEP
                chrom, start, step, span = self._parse_fixed_step_header(line)                
                if self.start is None:
                    self.start = start
                elif ((self.chrom != chrom) or
                      ((start - self.start) > chunk_size)):
                    yield self.start - 1, self.pos - 1, arr[:self.pos - self.start] # yield current chunk
                    arr[:] = 0 # reset array
                    self.start = start
                self.chrom, self.step, self.span = chrom, step, span
                self.pos = self.start                                
            elif line.startswith('variableStep'):
                self.state = self.VARIABLE_STEP
                chrom, span = self._parse_variable_step_header(line)                
                if (self.chrom != chrom):
                    yield self.start - 1, self.pos - 1, arr[:self.pos - self.start] # yield current chunk
                    arr[:] = 0 # reset array
                self.chrom, self.span = chrom, span
            else:                                        
                if self.state == self.FIXED_STEP:                    
                    # check if need to reset the chunk
                    if (self.pos + self.span - self.start) > chunk_size:
                        yield self.start - 1, self.pos - 1, arr[:self.pos - self.start] # yield current chunk
                        arr[:] = 0 # reset array
                        self.start = self.pos
                    arr[self.pos - self.start:self.pos + self.span - self.start] = self.dtype(line)
                    self.pos += self.step
                elif self.state == self.VARIABLE_STEP:
                    fields = line.split('\t')
                    pos, value = int(fields[0]), self.dtype(fields[1]) 
                    # check if need to reset the chunk
                    if (pos + self.span - self.start) > chunk_size:
                        yield self.start - 1, self.pos - 1, arr[:self.pos - self.start] # yield current chunk
                        arr[:] = 0 # reset array
                        self.start = self.pos
                    self.pos = pos
                    arr[self.pos - self.start:self.pos + self.span - self.start] = value
        # check if need to return the last chunk
        if self.pos > self.start:
            yield self.start - 1, self.pos - 1, arr[:self.pos - self.start] # yield current chunk
                    
                    
                    
                    